from volprims_linear import *

def main():
    '''
    Set up a random scene, a camera, and render the scene.
    '''
    image_width = 128
    image_height = 128
    num_primitives = 20

    # 1. Create and initialize primitives
    # primitives = UniformVolumePrimitives(num_spheres=num_primitives)
    primitives = UniformVolumePrimitives(num_spheres=num_primitives)
    primitives.random_initialization(device=default_device)

    primitives.centers[0] = torch.tensor([1.0, 0.0, 0.0], device=default_device)  # Center the primitive at the origin
    primitives.scales[0] = torch.tensor([2.0, 4.0, 7.0], device=default_device)  # Set a uniform scale
    primitives.opacities[0] = torch.tensor([0.1], device=default_device)  # Set a moderate opacity

    # 2. Set up the camera for a 45-degree downward looking view
    # Camera positioned along the y-z plane, looking towards origin
    # A position like (0, -d, d) looking at (0,0,0) gives a 45-degree downward pitch.
    cam_pos_y = -world_radius * 1.5
    cam_pos_z = world_radius * 1.5
    
    camera_position = torch.tensor([0.0, cam_pos_y, cam_pos_z], device=default_device)
    look_at_point = torch.tensor([0.0, 0.0, 0.0], device=default_device)
    # Up vector (0,1,1) ensures the camera is tilted correctly.
    # The camera's "local up" will be derived from this and the forward vector.
    up_vector = torch.tensor([0.0, 1.0, 1.0], device=default_device) 
    fov_degrees = 60.0

    camera = PinHoleCamera(
        position=camera_position,
        look_at=look_at_point,
        up=up_vector,
        fov=fov_degrees,
        device=default_device
    )

    # 3. Render the scene
    print(f"Rendering a {image_width}x{image_height} image of {num_primitives} primitives...")
    rendered_image = render_scene(camera, primitives, image_width, image_height)

    # Normalize color
    rendered_image[:, :, :3] = rendered_image[:, :, :3] / torch.clamp(rendered_image[:, :, 3:4], min=1e-6)

    print(f"Rendered image shape: {rendered_image.shape}")
    
    # Display the image using matplotlib
    import matplotlib.pyplot as plt
    # Ensure the image is on the CPU and converted to a NumPy array.
    # The alpha channel is the 4th channel. imshow handles RGBA automatically.
    image_to_show = rendered_image.cpu().numpy()
    
    plt.imshow(image_to_show)
    plt.title("Rendered Scene")
    plt.show()
    
    # Optionally, save the image
    # plt.imsave("rendered_scene.png", image_to_show)

if __name__ == "__main__":
    main()



